# color_clustering_batch.py
# -*- coding: utf-8 -*-
import cv2
import numpy as np
from sklearn.cluster import KMeans
import json
import os
import shutil

# -----------------------------
# 参数区（确保路径正确）
# -----------------------------
INPUT_DIR = r"D:\livan\codes\TongueClassify\Data\第二批_备份_提取整理"  # 原始字符串避免转义
OUTPUT_DIR = r"D:\livan\codes\TongueClassify\Data\output"
NUM_CLUSTERS = 5
MOVE_BY_DOMINANT = True
IMG_EXTS = {'.jpg', '.jpeg', '.png', '.bmp'}

# 颜色映射（保持不变）
REF_COLORS_LAB = np.array([[95, -2, 2], [100, 0, 0], [40, 60, -20], [50, 65, 35], [80, -5, 70]], dtype=np.float32)
REF_LABELS = ["淡白", "白", "紫暗", "红", "黄"]
DISPLAY_COLORS = {"淡白": (245, 245, 245), "白": (255, 255, 255), "紫暗": (128, 0, 128), "红": (255, 0, 0),
                  "黄": (255, 255, 0)}
DOMINANT_DIR = {"淡白": "light_white", "白": "white", "紫暗": "dark_purple", "红": "red", "黄": "yellow"}


def ensure_dir(path):
    os.makedirs(path, exist_ok=True)


def process_image(img_path, output_dir, num_clusters):
    img_path = os.path.normpath(img_path)

    # 1. 验证文件存在且可读
    if not os.path.isfile(img_path):
        print(f"错误：文件不存在 -> {img_path}")
        return

    # 2. 尝试读取图像（优先OpenCV，失败则用Pillow）
    try:
        img_bgr = cv2.imread(img_path)
        if img_bgr is None:
            raise ValueError("OpenCV读取失败，尝试Pillow")

        # 检查图像是否为空（可能损坏）
        if img_bgr.size == 0:
            raise ValueError("图像数据为空")

    except Exception as e:
        try:
            from PIL import Image
            img_pil = Image.open(img_path).convert("RGB")
            img_bgr = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
        except Exception as e:
            print(f"致命错误：无法读取 {img_path} -> {str(e)}")
            return

    # 3. 颜色聚类（后续代码不变）
    img_lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
    h, w = img_lab.shape[:2]
    pixels = img_lab.reshape(-1, 3)

    kmeans = KMeans(n_clusters=num_clusters, random_state=42)
    labels = kmeans.fit_predict(pixels)
    centers = kmeans.cluster_centers_

    dists = np.linalg.norm(centers[:, None, :] - REF_COLORS_LAB[None, :, :], axis=2)
    center_to_ref = np.argmin(dists, axis=1)
    center_labels = [REF_LABELS[i] for i in center_to_ref]
    semantic_map = np.array([center_labels[l] for l in labels]).reshape(h, w)

    # 4. 保存可视化结果和JSON
    base = os.path.splitext(os.path.basename(img_path))[0]
    vis_path = os.path.join(output_dir, f"semantic_{base}.png")
    json_path = os.path.join(output_dir, f"color_clusters_{base}.json")

    vis = np.zeros((h, w, 3), dtype=np.uint8)
    for lbl, col in DISPLAY_COLORS.items():
        vis[semantic_map == lbl] = col
    cv2.imwrite(vis_path, vis)

    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump({"reference_labels": REF_LABELS, "semantic_map": semantic_map.tolist()}, f, ensure_ascii=False,
                  indent=4)
    print(f"处理完成：{base}")

    # 5. 按主色移动文件
    if MOVE_BY_DOMINANT:
        counts = {lbl: int((semantic_map == lbl).sum()) for lbl in REF_LABELS}
        dominant = max(counts, key=counts.get)
        subdir = os.path.join(output_dir, DOMINANT_DIR[dominant])
        ensure_dir(subdir)
        shutil.copy(img_path, os.path.join(subdir, os.path.basename(img_path)))
        print(f"  主色分类：{dominant} -> {subdir}")


def batch_process(input_dir, output_dir, num_clusters):
    ensure_dir(output_dir)
    if MOVE_BY_DOMINANT:
        for d in DOMINANT_DIR.values():
            ensure_dir(os.path.join(output_dir, d))

    # 遍历文件并过滤有效图像
    for fn in os.listdir(input_dir):
        ext = os.path.splitext(fn)[1].lower()
        if ext in IMG_EXTS:
            file_path = os.path.join(input_dir, fn)
            process_image(file_path, output_dir, num_clusters)


if __name__ == "__main__":
    print(f"开始处理：输入目录={INPUT_DIR}, 输出目录={OUTPUT_DIR}")
    batch_process(INPUT_DIR, OUTPUT_DIR, NUM_CLUSTERS)
    print("处理完成！")